import torch
import torch.nn as nn

class Encoder(nn.Module):
    def __init__(self,encoded_space_dim,fc2_input_dim):
        super(Encoder, self).__init__()
        #构建卷积网络
        self.encoder_cnn = nn.Sequential(
            nn.Conv2d(1,8,3,stride=2,padding=1),
            nn.ReLU(),
            # second cnn
            nn.Conv2d(8,16,3,stride=2,padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            # three cnn
            nn.Conv2d(16,32,3,stride=2,padding=0),
            nn.ReLU()
        )

        self.flatten = nn.Flatten(start_dim=1)
        #全连接网络
        self.encoder_line = nn.Sequential(
            # first linear
            nn.Linear(3*3*32,128),
            nn.Linear(128,encoded_space_dim)
        )

    def forward(self,x):
        x = self.encoder_cnn(x)
        x = self.flatten(x)
        x = self.encoder_line(x)
        return x

